use slab::Slab;
use std::fmt::{self, Debug};

#[cfg(test)]
mod test;

#[derive(Debug, Clone, Copy, PartialEq)]
enum Color {
    Red,
    Black,
}

#[derive(Debug, Clone)]
struct RBTNode<K, V> {
    color: Color,
    key: K,
    value: V,
    left: Option<usize>,
    right: Option<usize>,
    parent: Option<usize>,
}

impl<K, V> RBTNode<K, V> {
    fn new(key: K, value: V) -> Self {
        Self {
            color: Color::Red,
            key: key,
            value: value,
            left: None,
            right: None,
            parent: None,
        }
    }
}

pub struct RBTree<K, V> {
    memory: Slab<RBTNode<K, V>>,
    root: Option<usize>,
}

impl<K: Ord + Debug + Clone, V: Debug + Clone> RBTree<K, V> {
    pub fn new() -> Self {
        Self {
            memory: Slab::default(),
            root: None,
        }
    }

    pub fn insert(&mut self, key: K, value: V) {
        let new_node = self.memory.insert(RBTNode::new(key.clone(), value));
        let mut y = None;
        let mut x = self.root;

        while let Some(x_ref) = x {
            y = Some(x_ref);
            let x_node = self.memory.get(x_ref).unwrap();
            if key < x_node.key {
                x = x_node.left;
            } else {
                x = x_node.right;
            }
        }

        let z = self.memory.get_mut(new_node).unwrap();
        z.parent = y;
        match y {
            None => self.root = Some(new_node),
            Some(y_ref) => {
                let y_node = self.memory.get_mut(y_ref).unwrap();
                if key < y_node.key {
                    y_node.left = Some(new_node);
                } else {
                    y_node.right = Some(new_node);
                }
            }
        }

        self.insert_fixup(new_node);
    }

    fn insert_fixup(&mut self, mut z: usize) {
        while let Some(z_parent) = self.memory.get(z).unwrap().parent {
            if self.memory.get(z_parent).unwrap().color == Color::Black {
                break;
            }
            let z_grandparent = self.memory.get(z_parent).unwrap().parent.expect("空指针");
            if Some(z_parent) == self.memory.get(z_grandparent).unwrap().left {
                let y = self.memory.get(z_grandparent).unwrap().right;
                if let Some(y_ref) = y {
                    if self.memory.get(y_ref).unwrap().color == Color::Red {
                        self.memory.get_mut(z_parent).unwrap().color = Color::Black;
                        self.memory.get_mut(y_ref).unwrap().color = Color::Black;
                        self.memory.get_mut(z_grandparent).unwrap().color = Color::Red;
                        z = z_grandparent;
                        continue;
                    }
                }
                if Some(z) == self.memory.get(z_parent).unwrap().right {
                    z = z_parent;
                    self.left_rotate(z);
                }
                let z_parent = self.memory.get(z).unwrap().parent.expect("空指针");
                let z_grandparent = self.memory.get(z_parent).unwrap().parent.expect("空指针");
                self.memory.get_mut(z_parent).unwrap().color = Color::Black;
                self.memory.get_mut(z_grandparent).unwrap().color = Color::Red;
                self.right_rotate(z_grandparent);
            } else {
                // 对称情况，左右互换
                let y = self.memory.get(z_grandparent).unwrap().left;
                if let Some(y_ref) = y {
                    if self.memory.get(y_ref).unwrap().color == Color::Red {
                        self.memory.get_mut(z_parent).unwrap().color = Color::Black;
                        self.memory.get_mut(y_ref).unwrap().color = Color::Black;
                        self.memory.get_mut(z_grandparent).unwrap().color = Color::Red;
                        z = z_grandparent;
                        continue;
                    }
                }
                if Some(z) == self.memory.get(z_parent).unwrap().left {
                    z = z_parent;
                    self.right_rotate(z);
                }
                let z_parent = self.memory.get(z).unwrap().parent.expect("空指针");
                let z_grandparent = self.memory.get(z_parent).unwrap().parent.expect("空指针");
                self.memory.get_mut(z_parent).unwrap().color = Color::Black;
                self.memory.get_mut(z_grandparent).unwrap().color = Color::Red;
                self.left_rotate(z_grandparent);
            }
        }
        if let Some(root) = self.root {
            self.memory.get_mut(root).unwrap().color = Color::Black;
        }
    }

    fn left_rotate(&mut self, x: usize) {
        let y = self.memory.get(x).unwrap().right.expect("右子节点不能为空");
        self.memory.get_mut(x).unwrap().right = self.memory.get(y).unwrap().left;
        
        if let Some(y_left) = self.memory.get(y).unwrap().left {
            self.memory.get_mut(y_left).unwrap().parent = Some(x);
        }
        
        let x_parent = self.memory.get(x).unwrap().parent;
        self.memory.get_mut(y).unwrap().parent = x_parent;
        
        match x_parent {
            None => self.root = Some(y),
            Some(parent) => {
                let parent_node = self.memory.get_mut(parent).unwrap();
                if Some(x) == parent_node.left {
                    parent_node.left = Some(y);
                } else {
                    parent_node.right = Some(y);
                }
            }
        }
        
        self.memory.get_mut(y).unwrap().left = Some(x);
        self.memory.get_mut(x).unwrap().parent = Some(y);
    }

    fn right_rotate(&mut self, y: usize) {
        let x = self.memory.get(y).unwrap().left.expect("左子节点不能为空");
        self.memory.get_mut(y).unwrap().left = self.memory.get(x).unwrap().right;
        
        if let Some(x_right) = self.memory.get(x).unwrap().right {
            self.memory.get_mut(x_right).unwrap().parent = Some(y);
        }
        
        let y_parent = self.memory.get(y).unwrap().parent;
        self.memory.get_mut(x).unwrap().parent = y_parent;
        
        match y_parent {
            None => self.root = Some(x),
            Some(parent) => {
                let parent_node = self.memory.get_mut(parent).unwrap();
                if Some(y) == parent_node.right {
                    parent_node.right = Some(x);
                } else {
                    parent_node.left = Some(x);
                }
            }
        }
        
        self.memory.get_mut(x).unwrap().right = Some(y);
        self.memory.get_mut(y).unwrap().parent = Some(x);
    }

    #[allow(unused)]
    fn inorder_helper(&self, node: Option<usize>) {
        if let Some(node_ref) = node {
            let node = self.memory.get(node_ref).unwrap();
            self.inorder_helper(node.left);
            println!(
                "(color: {:?}, key: {:?}, value: {:?})",
                node.color, node.key, node.value
            );
            self.inorder_helper(node.right);
        }
    }

    #[allow(unused)]
    fn inorder(&self) {
        self.inorder_helper(self.root);
    }

    pub fn remove(&mut self, key: &K) {
        if let Some(z) = self.find_node(key) {
            let z_color = self.memory.get(z).unwrap().color;
            let mut x;
            let mut x_parent;

            if self.memory.get(z).unwrap().left.is_none() {
                x = self.memory.get(z).unwrap().right;
                x_parent = self.memory.get(z).unwrap().parent;
                self.transplant(z, self.memory.get(z).unwrap().right);
            } else if self.memory.get(z).unwrap().right.is_none() {
                x = self.memory.get(z).unwrap().left;
                x_parent = self.memory.get(z).unwrap().parent;
                self.transplant(z, self.memory.get(z).unwrap().left);
            } else {
                let y = self.minimum(self.memory.get(z).unwrap().right.unwrap());
                let y_color = self.memory.get(y).unwrap().color;
                x = self.memory.get(y).unwrap().right;
                x_parent = Some(y);

                if self.memory.get(y).unwrap().parent == Some(z) {
                    if let Some(x_ref) = x {
                        self.memory.get_mut(x_ref).unwrap().parent = Some(y);
                    }
                } else {
                    x_parent = self.memory.get(y).unwrap().parent;
                    self.transplant(y, self.memory.get(y).unwrap().right);
                    let z_right = self.memory.get(z).unwrap().right.unwrap();
                    self.memory.get_mut(y).unwrap().right = Some(z_right);
                    self.memory.get_mut(z_right).unwrap().parent = Some(y);
                }

                self.transplant(z, Some(y));
                let z_left = self.memory.get(z).unwrap().left.unwrap();
                self.memory.get_mut(y).unwrap().left = Some(z_left);
                self.memory.get_mut(z_left).unwrap().parent = Some(y);
                self.memory.get_mut(y).unwrap().color = self.memory.get(z).unwrap().color;
            }

            self.memory.remove(z);

            if z_color == Color::Black {
                self.remove_fixup(x, x_parent);
            }
        }
    }

    fn remove_fixup(&mut self, mut x: Option<usize>, mut x_parent: Option<usize>) {
        while x != self.root && x.map_or(true, |node| self.memory.get(node).unwrap().color == Color::Black) {
            if x == x_parent.and_then(|p| self.memory.get(p).unwrap().left) {
                let mut w = x_parent.and_then(|p| self.memory.get(p).unwrap().right);
                if let Some(w_ref) = w {
                    if self.memory.get(w_ref).unwrap().color == Color::Red {
                        self.memory.get_mut(w_ref).unwrap().color = Color::Black;
                        self.memory.get_mut(x_parent.unwrap()).unwrap().color = Color::Red;
                        self.left_rotate(x_parent.unwrap());
                        w = self.memory.get(x_parent.unwrap()).unwrap().right;
                    }
                    if w.and_then(|w| self.memory.get(w).unwrap().left)
                        .map_or(true, |left| self.memory.get(left).unwrap().color == Color::Black)
                        && w.and_then(|w| self.memory.get(w).unwrap().right)
                        .map_or(true, |right| self.memory.get(right).unwrap().color == Color::Black)
                    {
                        self.memory.get_mut(w.unwrap()).unwrap().color = Color::Red;
                        x = x_parent;
                        x_parent = x.and_then(|x| self.memory.get(x).unwrap().parent);
                    } else {
                        if w.and_then(|w| self.memory.get(w).unwrap().right)
                            .map_or(true, |right| self.memory.get(right).unwrap().color == Color::Black)
                        {
                            w.and_then(|w| self.memory.get(w).unwrap().left)
                                .map(|left| self.memory.get_mut(left).unwrap().color = Color::Black);
                            self.memory.get_mut(w.unwrap()).unwrap().color = Color::Red;
                            self.right_rotate(w.unwrap());
                            w = self.memory.get(x_parent.unwrap()).unwrap().right;
                        }
                        let x_parent_color = self.memory.get(x_parent.unwrap()).unwrap().color;
                        self.memory.get_mut(w.unwrap()).unwrap().color = x_parent_color;
                        self.memory.get_mut(x_parent.unwrap()).unwrap().color = Color::Black;
                        w.and_then(|w| self.memory.get(w).unwrap().right)
                            .map(|right| self.memory.get_mut(right).unwrap().color = Color::Black);
                        self.left_rotate(x_parent.unwrap());
                        x = self.root;
                        break;
                    }
                } else {
                    x = x_parent;
                    x_parent = x.and_then(|x| self.memory.get(x).unwrap().parent);
                }
            } else {
                // 对称情况，左右互换
                let mut w = x_parent.and_then(|p| self.memory.get(p).unwrap().left);
                if let Some(w_ref) = w {
                    if self.memory.get(w_ref).unwrap().color == Color::Red {
                        self.memory.get_mut(w_ref).unwrap().color = Color::Black;
                        self.memory.get_mut(x_parent.unwrap()).unwrap().color = Color::Red;
                        self.right_rotate(x_parent.unwrap());
                        w = self.memory.get(x_parent.unwrap()).unwrap().left;
                    }
                    if w.and_then(|w| self.memory.get(w).unwrap().right)
                        .map_or(true, |right| self.memory.get(right).unwrap().color == Color::Black)
                        && w.and_then(|w| self.memory.get(w).unwrap().left)
                        .map_or(true, |left| self.memory.get(left).unwrap().color == Color::Black)
                    {
                        self.memory.get_mut(w.unwrap()).unwrap().color = Color::Red;
                        x = x_parent;
                        x_parent = x.and_then(|x| self.memory.get(x).unwrap().parent);
                    } else {
                        if w.and_then(|w| self.memory.get(w).unwrap().left)
                            .map_or(true, |left| self.memory.get(left).unwrap().color == Color::Black)
                        {
                            w.and_then(|w| self.memory.get(w).unwrap().right)
                                .map(|right| self.memory.get_mut(right).unwrap().color = Color::Black);
                            self.memory.get_mut(w.unwrap()).unwrap().color = Color::Red;
                            self.left_rotate(w.unwrap());
                            w = self.memory.get(x_parent.unwrap()).unwrap().left;
                        }
                        let x_parent_color = self.memory.get(x_parent.unwrap()).unwrap().color;
                        self.memory.get_mut(w.unwrap()).unwrap().color = x_parent_color;
                        self.memory.get_mut(x_parent.unwrap()).unwrap().color = Color::Black;
                        w.and_then(|w| self.memory.get(w).unwrap().left)
                            .map(|left| self.memory.get_mut(left).unwrap().color = Color::Black);
                        self.right_rotate(x_parent.unwrap());
                        x = self.root;
                        break;
                    }
                } else {
                    x = x_parent;
                    x_parent = x.and_then(|x| self.memory.get(x).unwrap().parent);
                }
            }
        }
        if let Some(x_ref) = x {
            self.memory.get_mut(x_ref).unwrap().color = Color::Black;
        }
    }

    pub fn len(&self) -> usize {
        self.memory.len()
    }


    fn transplant(&mut self, u: usize, v: Option<usize>) {
        let u_parent = self.memory.get(u).unwrap().parent;
        match u_parent {
            None => self.root = v,
            Some(parent) => {
                let parent_node = self.memory.get_mut(parent).unwrap();
                if Some(u) == parent_node.left {
                    parent_node.left = v;
                } else {
                    parent_node.right = v;
                }
            }
        }
        if let Some(v_ref) = v {
            self.memory.get_mut(v_ref).unwrap().parent = u_parent;
        }
    }

    fn minimum(&self, mut x: usize) -> usize {
        while let Some(left) = self.memory.get(x).unwrap().left {
            x = left;
        }
        x
    }

    fn find_node(&self, key: &K) -> Option<usize> {
        let mut current = self.root;
        while let Some(node) = current {
            let node_ref = self.memory.get(node).unwrap();
            match key.cmp(&node_ref.key) {
                std::cmp::Ordering::Equal => return Some(node),
                std::cmp::Ordering::Less => current = node_ref.left,
                std::cmp::Ordering::Greater => current = node_ref.right,
            }
        }
        None
    }

    pub fn modify(&mut self, key: &K, value: V) {
        if let Some(node) = self.find_node(key) {
            let node_ref = self.memory.get_mut(node).unwrap();
            node_ref.value = value;
        }
    }

    pub fn iter(&self) -> RBTreeIterator<K, V> {
        self.into_iter()
    }

    pub fn get(&self, key: &K) -> Option<&V> {
        self.find_node(key).map(|node| &self.memory.get(node).unwrap().value)
    }


    #[allow(unused)]
    fn is_balanced(&self) -> bool {
        if self.root.is_none() {
            return true;
        }
        if self.memory.get(self.root.expect("nullptr")).unwrap().color != Color::Black {
            println!("根节点不是黑色");
            return false;
        }
        let mut benchmark = 0;
        let mut cur = self.root;
        while let Some(node) = cur {
            if self.memory.get(node).unwrap().color == Color::Black {
                benchmark += 1;
            }
            cur = self.memory.get(node).unwrap().left;
        }

        return self.is_balanced_helper(self.root, 0, &mut benchmark);
    }

    #[allow(unused)]
    fn is_balanced_helper(&self, node: Option<usize>, black_num: i32, benchmark: &mut i32) -> bool {
        let mut black_num = black_num;
        if node.is_none() {
            if *benchmark == 0 {
                *benchmark = black_num;
                return true;
            }
            if black_num != *benchmark {
                println!("某条黑色结点数量不相等");
                return false;
            } else {
                return true;
            }
        }
        if self.memory.get(node.expect("nullptr")).unwrap().color == Color::Black {
            black_num += 1;
        }
        if self.memory.get(node.expect("nullptr")).unwrap().color == Color::Red && self.memory.get(self.memory.get(node.expect("nullptr")).unwrap().parent.expect("nullptr")).unwrap().color == Color::Red {
            println!("出现连续红色结点");
            return false;
        }
        self.is_balanced_helper(self.memory.get(node.expect("nullptr")).unwrap().left, black_num, benchmark) && self.is_balanced_helper(self.memory.get(node.expect("nullptr")).unwrap().right, black_num, benchmark)
    }
}

impl fmt::Display for Color {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Color::Red => write!(f, "Red"),
            Color::Black => write!(f, "Black"),
        }
    }
}

// 定义迭代器结构
pub struct RBTreeIterator<'a, K, V> {
    tree: &'a RBTree<K, V>,
    current: Option<usize>,
}

// 为迭代器实现 Iterator trait
impl<'a, K: Ord + Debug + Clone, V: Debug + Clone> Iterator for RBTreeIterator<'a, K, V> {
    type Item = (&'a K, &'a V);

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(node) = self.current {
            let node_ref = self.tree.memory.get(node).unwrap();
            let result = Some((&node_ref.key, &node_ref.value));
            self.current = self.tree.next_node(node);
            result
        } else {
            None
        }
    }
}

// 为 RBTree 实 IntoIterator trait
impl<'a, K: Ord + Debug + Clone, V: Debug + Clone> IntoIterator for &'a RBTree<K, V> {
    type Item = (&'a K, &'a V);
    type IntoIter = RBTreeIterator<'a, K, V>;

    fn into_iter(self) -> Self::IntoIter {
        RBTreeIterator {
            tree: self,
            current: self.root.and_then(|root| self.leftmost_node(root)),
        }
    }
}

// 在 RBTree 实现中添加以下辅助方法
impl<K: Ord + Debug + Clone, V: Debug + Clone> RBTree<K, V> {
    fn leftmost_node(&self, node: usize) -> Option<usize> {
        let mut current = node;
        while let Some(left) = self.memory.get(current).unwrap().left {
            current = left;
        }
        Some(current)
    }

    fn next_node(&self, node: usize) -> Option<usize> {
        let node_ref = self.memory.get(node).unwrap();
        if let Some(right) = node_ref.right {
            self.leftmost_node(right)
        } else {
            self.find_successor(node)
        }
    }

    fn find_successor(&self, node: usize) -> Option<usize> {
        let mut current = node;
        let mut parent = self.memory.get(current).unwrap().parent;

        while let Some(p) = parent {
            if self.memory.get(p).unwrap().left == Some(current) {
                return Some(p);
            }
            current = p;
            parent = self.memory.get(current).unwrap().parent;
        }

        None
    }
}